# !/usr/bin/python2.7
# -*- coding: UTF-8 -*-

import os
import sys
import time
import datetime
import binascii
import struct
import mblue
import mcurses
from ctypes import *

class FloatBits(Union):
    class _Bits(Structure):
        _fields_ = [('M',c_uint, 4),
                   ('E',c_uint, 9),
                  ('S',c_uint, 19) ]
    _anonymous_ = ('bits', )
    _fields_ = [('value',c_uint),
               ('bits', _Bits) ]

class ota_file_header(Structure):
    _pack_ = 1
    _fields_ = [('headerver',c_uint32),         #头格式
                ('devtype',c_uint32),           #机型,0-无效,1-spu100,2-spu200
                ('object', c_uint32),           #作用对象,0-无效,1-efr32,2-stm32,3-hr
                ('prgmver',c_uint32),           #程序版本 xx.xx.xx.xx
                ('prgmsize',c_uint32),          #固件大小(字节)
                ('prgmoffset',c_uint32),        #固件偏移地址
                ('checkcd',c_uint32)]           #程序固件校验码CRC32,不包括该头

class ota_header(Structure):
    _pack_ = 1
    _fields_ = [('type_hilen',c_uint8),     #低3位为数据段长度高3位
                ('lolen',c_uint8),          #数据段长度低8位
                ('cls',c_uint8),            #0x00-upload 类
                ('method',c_uint8),         #命令类型,0x20-Message ID 
                ('address',c_uint32)]       #该项在升级时指明每包固件的存放偏移地址

class mt_ota_header(Structure):
    _pack_ = 1
    _fields_ = [('begin',c_uint16),
                ('pkgtype',c_uint8),
                ('sn',c_uint8),
                ('dlen',c_uint8),
                ('datas',c_uint8*256)]

def dbgprint(istr,win = None):
    if win == None:
        print istr
    else:
        win.mprint(istr + '\n')

def ota_peek_event(dev,chand,timeout = 3):
    read=dev.read_char(chand,timeout)
    if read is not None:
        read=read['d'][0]
        try:
            read.decode('utf-8')
        except UnicodeDecodeError as e:
            #read=binascii.b2a_hex(read)
            #print binascii.b2a_hex(read)
            if len(read) < 10:
                return None
            if read[0:2] == '\xBE\xBE' and read[-2:] == '\xED\xED':
                read = cast(read,POINTER(mt_ota_header))
                #print '%d'%(read.contents.pkgtype)
                if read.contents.pkgtype == 0x31:
                    read = bytearray(read.contents.datas)
                    #print binascii.b2a_hex(read)
                else:
                    return None
            else:
                return None
            #i = read.find('\x20')
            i = 0
            if i != (-1) and len(read[i:]) > 7:
                header = ota_header.from_buffer(bytearray(read[i:]))
                msg = (header.method << 24) + (header.cls << 16) + (header.lolen << 8) + header.type_hilen
                return header,read[i:][i+8:i+8+((header.type_hilen & 0x07) << 8) + header.lolen],(msg & 0xffff00f8)
    return None

def file_effective(filepath):
    return (os.path.exists(filepath) and (os.path.getsize(filepath)) > 0)

def sync_dev(dev,chand):
    otaheader = ota_header(0x20,0x01,0x00,0x00,0x00)
    mtotaheader = mt_ota_header()
    mtotaheader.begin = 0xBEBE
    mtotaheader.pkgtype = 0x31
    mtotaheader.sn = 0x00
    mtotaheader.dlen = 0x09
    mtotaheader.datas = tuple(bytearray(otaheader) + '\x00' + '\xED\xED') #(otaheader.type_hilen,otaheader.lolen,otaheader.cls,otaheader.method,otaheader.address,0x00,0xED,0xED)#倒数第三个字节，这里是0x02，代表要升级蓝牙程序,0x10代表要升级STM32程序,0x11代表要升级心率表程序,0x00由固件指明此处忽略
    while True:
        #dbgprint('.',win)
        dev.write_char(chand,bytearray(mtotaheader)[0:16],True)
        #print binascii.b2a_hex(bytearray(mtotaheader)[0:16])
        rsp = ota_peek_event(dev,chand)
        if (rsp != None) and (rsp[2] == 0x000000a0):
            #dbgprint('sync ok',win)
            dbgprint('Current program version: v%d.%d.%d'%(rsp[1][2],rsp[1][1],rsp[1][0]))
            #clear reciv buff
            return True
        #time.sleep(1)
        mtotaheader.sn = mtotaheader.sn + 1
    return False

def ota_file_data_upload(dev,chand,size,data,address):
    if size > 230:
        return 'data size too large'
    if size <= 0:
        return 'no data,how send?'
    otaheader = ota_header(0x20,0x00,0x00,0x02,0x00)
    otaheader.type_hilen = (0x20 | ((size >> 8) & 0x07)) 
    otaheader.lolen = (size & 0xff)
    otaheader.address = address
    mtotaheader = mt_ota_header()
    mtotaheader.begin = 0xBEBE
    mtotaheader.pkgtype = 0x31
    mtotaheader.sn = 0x00
    mtotaheader.dlen = (size + 8)
    mtotaheader.datas = tuple(bytearray(otaheader) + bytearray(data)[0:size] + '\xED\xED') 

    retrycount = 16
    timeout = 1
    while retrycount:
        dev.write_char(chand,bytearray(mtotaheader)[0:(13 + size + 2)],True)
        #print '\r\nsend: %08X'%(address-128),binascii.b2a_hex(bytearray(mtotaheader)[0:(13 + size + 2)])
        rsp = ota_peek_event(dev,chand,timeout)
        #print 'rsp:',binascii.b2a_hex(bytearray(rsp[0])),binascii.b2a_hex(rsp[1]),hex(rsp[2])
            #time.sleep(0.2)
            #return True
        #if (rsp != None) and (rsp[2] == 0x02000020) and (((rsp[1][1] << 8) + rsp[1][0]) == size):
        #    return True
        #elif rsp == None:
        #    return '\nno answer'
        #else:
        #    return ('\nanswer error,receive error (%#X,0x%04x,0x%04x)' % (rsp[2],((rsp[1][1] << 8) + rsp[1][0]),size))
    
        if (rsp != None) and (rsp[2] == 0x02000020) and (rsp[0].address == address) and (((rsp[1][1] << 8) + rsp[1][0]) == size):
            return True
        else:
            retrycount = retrycount - 1
            mtotaheader.sn = mtotaheader.sn + 1
            timeout = 1

    if rsp == None:
        return '\nno answer'
    else:
        return ('\nanswer error,receive error (%#X,0x%04x,0x%04x,0x%08X,0x%08X)' % (rsp[2],((rsp[1][1] << 8) + rsp[1][0]),size,rsp[0].address,address))

def ota_upload_finish(win,dev,chand):
    otaheader = ota_header(0x20,0x00,0x00,0x03,0x00)
    mtotaheader = mt_ota_header()
    mtotaheader.begin = 0xBEBE
    mtotaheader.pkgtype = 0x31
    mtotaheader.sn = 0x00
    mtotaheader.dlen = 0x08
    mtotaheader.datas = tuple(bytearray(otaheader) + '\xED\xED')#(otaheader.type_hilen,otaheader.lolen,otaheader.cls,otaheader.method,otaheader.address,0xED,0xED)

    i = 0 
    getrsp = 0
    while (i < 300):
        if getrsp == 0:
            dev.write_char(chand,bytearray(mtotaheader)[0:15],True)
        rsp = ota_peek_event(dev,chand)
        if (rsp != None) and (rsp[2] == 0x03000020):
            getrsp = 1
            if (((rsp[1][1] << 8) + rsp[1][0]) == 0x0001):#接收完成，开始校验
                dbgprint('receive complete,start verify...',win)
            elif (((rsp[1][1] << 8) + rsp[1][0]) == 0x0002):#校验成功，开始升级
                dbgprint('verify complete,start upgrade...',win)
            elif (((rsp[1][1] << 8) + rsp[1][0]) == 0x0003):#升级成功
                dbgprint('upgrade success',win)
                return True
            elif (((rsp[1][1] << 8) + rsp[1][0]) == 0x0101):#接收错误
                dbgprint('receive error',win)
                return False
            elif (((rsp[1][1] << 8) + rsp[1][0]) == 0x0102):#校验错误
                dbgprint('verify error',win)
                return False
            elif (((rsp[1][1] << 8) + rsp[1][0]) == 0x0103):#升级失败
                dbgprint('upgrade fail',win)
                return False
        else:
            i = i + 1    
    dbgprint('time out',win)
    return False

def upload(win,dev,chand,filepath):
    readsize = 224
    address = 0
    filesize = os.path.getsize(filepath)
    if win != None:
        (y,x) = win.getscr().getyx()
    bgtime = datetime.datetime.now()
    dbgprint('send pkg size: %d + 15'%(readsize),win)
    with open(filepath,'r') as fs:
        ota_data = fs.read(readsize) #128
        ota_size = len(ota_data)
        while ota_size != 0:
            rsp = ota_file_data_upload(dev,chand,ota_size,ota_data,address)
            address = address + ota_size
            t = (datetime.datetime.now()-bgtime).seconds
            if type(rsp) is str:
                dbgprint(rsp,win)
                return False
            ps = fs.tell() * 100.0 / filesize
            sp = (t > 0) and fs.tell()/(t) or fs.tell()
            if win == None:
                ps = int(ps)
                print '\r|%50s| [%3d%%] [%5d Bps]'%((ps/2) * '█' + (50 - ps/2) * ' ',ps,sp),
            else:
                win.progressbar(ps,0,y+1)
                win.display_info('[%5d Bps]'%(sp),60,y+1)

            ota_data = fs.read(readsize) #128
            ota_size = len(ota_data)

    dbgprint('\r\nUse time: %d s'%(t))
    return True

def main_loop(dev,chand,filepath,win = None):
    if not file_effective(filepath):
        dbgprint('file path error\r\n',win)
        return False
    filesize = os.path.getsize(filepath)
    dbgprint('file size: %d Byte\r\n' % (filesize),win)
    dbgprint('sync...',win)
    rsp = sync_dev(dev,chand)
    if rsp:
        dbgprint('sync ok,start upload',win)
    else:
        dbgprint('sync fail',win)
        return
    rsp = upload(win,dev,chand,filepath)
    if rsp:
        dbgprint('\r\nupload finish',win)
    else:
        dbgprint('\r\nupload fail',win)
        return
    ota_upload_finish(win,dev,chand)
    return

def printusage():
    print "Usage:"
    print "     mt_ota mac filepath"
    print "     mt_ota encry infile [outfile] [keyfile] [--devtype (1,2) --obj(1,2,3) --ver(x.x.x)]"
    return;

def mkupdatefirmware(infile,outfile,keyfile):
    crc = os.popen('./crc 32 %s'%(infile),'r').read()
    filesize = int(crc.split()[1])
    poly = crc.split()[2]
    crc = int(crc.split()[0][2:],16)
    offset = 128

    if '--devtype' not in sys.argv:
        print '请输入设备型号代码(1: SPU100,2: SPU200):',
        devnu = int(raw_input())
    else:
        devnu = int(sys.argv[sys.argv.index('--devtype') + 1]) 

    if '--obj' not in sys.argv:
        print '请输入作用对象(0: 无效,1: efr32,2: stm32,3: HR):',
        obj = int(raw_input())
    else:
        obj = int(sys.argv[sys.argv.index('--obj') + 1]) 

    if '--ver' not in sys.argv:
        print '请输入程序版本(x.x.x):',
        ver = raw_input().split('.') 
    else:
        ver = sys.argv[sys.argv.index('--ver') + 1].split('.')

    ver = int(ver[-1]) + (int(ver[-2]) << 8) + (int(ver[-3]) << 16) 
    print '你输入了:',devnu,obj,'%#X\r\n'%ver
    print '文件大小：',filesize
    print '生成校验码(CRC32 POLYNOMIAL=%s): %#X'%(poly,crc)
    
    print '添加文件头...'
    otafileheader = ota_file_header()
    otafileheader.headerver = 0x01      #头格式1 
    otafileheader.devtype = devnu       #机型,0-无效,1-spu100,2-spu200 
    otafileheader.object = obj          #作用对象,0-无效,1-efr32,2-stm32,3-HR
    otafileheader.prgmver = ver         #程序版本 xx.xx.xx.xx
    otafileheader.prgmsize = filesize   #固件大小(字节)
    otafileheader.prgmoffset = offset   #固件偏移地址 
    otafileheader.checkcd = crc         #程序固件校验码CRC32,不包括该头
    #print binascii.b2a_hex(bytearray(otafileheader)),len(bytearray(otafileheader)) 
    infilename = os.path.basename(infile)
    infilepath = infile.replace(infilename,'')
    with open(infile,'r') as fs:
        infile = infilepath + 'raw_' + infilename
        with open(infile,'w') as fo:
            lines = ''.join(fs.readlines())
            fo.write(bytearray(otafileheader))
            fo.write((otafileheader.prgmoffset - len(bytearray(otafileheader))) * '\x00')
            fo.write(lines)
    print '开始加密...'
    if outfile == None:
        outfile = infilepath + 'encry_' + '.'.join(infilename.split('.')[0:-1]) + '.ota'
    if keyfile == None:
        rsp = os.popen('./aes encry %s %s'%(infile,outfile),'r').read()
    else:
        rsp = os.popen('./aes encry %s %s %s'%(infile,outfile,keyfile),'r').read()
    print rsp
    print '加密文件：%s' % (outfile)
    os.system('rm %s'%infile)

    return

def main():
    if len(sys.argv) < 3:
        printusage()
        return
    firstagc = 0
    for agv in sys.argv:
        if '--' in agv:
            firstagc = sys.argv.index(agv)
            break
    if firstagc and (firstagc < 3):
        printusage()
        return
    if 'encry' in sys.argv:
        if len(sys.argv) == 3 or firstagc == 3:
            mkupdatefirmware(sys.argv[2],None,None)
        elif len(sys.argv) == 4 or firstagc == 4:
            mkupdatefirmware(sys.argv[2],sys.argv[3],None)
        elif len(sys.argv) == 5 or firstagc == 5:
            mkupdatefirmware(sys.argv[2],sys.argv[3],sys.argv[4])
        else:
            printusage()
            return

    else:
        mac = sys.argv[1]
        dev = mblue.bluedev()
        conhandle = dev.connect(mac)
        if not conhandle is None:
            print 'Connected'
            chand = dev.getcharhandleByUUID('1101','8888')
            if chand is not None:
                print 'discovery server'
                if dev.char_notify(chand) is True:
                    cnotify = True
                    print 'turn on notify success'
                else:
                    cnotify = False 
                    print 'turn on notify fail'
            else:
                cnotify = False 
                print 'no discovery server'
        main_loop(dev,chand,sys.argv[2])
        dev.disconnect()


if __name__ == '__main__':
    
    main()
    exit()




